#define CI_TILE 4
#define CO_TILE 4
#define UP_DIV(x, y) (((x) + (y) - (1)) / (y))

//#define CI 4
//#define IH 224
//#define IW 224
//#define CO 32
//#define OH 112
//#define OW 112
//#define KH 3
//#define KW 3
//#define strideH 2
//#define strideW 2
//#define padTop 0
//#define padLeft 0
//#define CI_SLICES 1
//#define CO_SLICES 8

//#define CI 960
//#define IH 7
//#define IW 7
//#define CO 320
//#define OH 7
//#define OW 7
//#define KH 1
//#define KW 1
//#define strideH 1
//#define strideW 1
//#define padTop 0
//#define padLeft 0
//#define CI_SLICES 240
//#define CO_SLICES 80

//inputNHWC_1x8x128x250_outputNHWC_1x8x128x100_kernelHW_1x1_strideHW_1x1_padTopBottomLeftRight_0x0x0x0_dilationHW_1x1
//#define CI 250
//#define IH 8
//#define IW 128
//#define CO 100
//#define OH 8
//#define OW 128
//#define KH 1
//#define KW 1
//#define strideH 1
//#define strideW 1
//#define padTop 0
//#define padLeft 0
//#define CI_SLICES 63
//#define CO_SLICES 25

// inputNHWC_1x36x14400x128_outputNHWC_1x36x14400x128_kernelHW_1x1_strideHW_1x1_padTopBottomLeftRight_0x0x0x0_dilationHW_1x1
//#define CI 128
//#define IH 36
//#define IW 14400
//#define CO 128
//#define OH 36
//#define OW 14400
//#define KH 1
//#define KW 1
//#define strideH 1
//#define strideW 1
//#define padTop 0
//#define padLeft 0
//#define CI_SLICES 32
//#define CO_SLICES 32

//inputNHWC_1x480x480x32_outputNHWC_1x480x480x128_kernelHW_1x1_strideHW_1x1_padTopBottomLeftRight_0x0x0x0_dilationHW_1x1
//#define CI 32
//#define IH 480
//#define IW 480
//#define CO 128
//#define OH 480
//#define OW 480
//#define KH 1
//#define KW 1
//#define strideH 1
//#define strideW 1
//#define padTop 0
//#define padLeft 0
//#define CI_SLICES 8
//#define CO_SLICES 32


//inputNHWC_1x32x512x50_outputNHWC_1x32x512x48_kernelHW_3x3_strideHW_1x1_padTopBottomLeftRight_1x1x1x1_dilationHW_1x1
//#define CI 50
//#define IH 32
//#define IW 512
//#define CO 48
//#define OH 32
//#define OW 512
//#define KH 3
//#define KW 3
//#define strideH 1
//#define strideW 1
//#define padTop 1
//#define padLeft 1
//#define CI_SLICES 13
//#define CO_SLICES 12

//inputNHWC_1x480x480x128_outputNHWC_1x480x480x128_kernelHW_3x3_strideHW_1x1_padTopBottomLeftRight_1x1x1x1_dilationHW_1x1
//#define CI 128
//#define IH 480
//#define IW 480
//#define CO 128
//#define OH 480
//#define OW 480
//#define KH 3
//#define KW 3
//#define strideH 1
//#define strideW 1
//#define padTop 1
//#define padLeft 1
//#define CI_SLICES 32
//#define CO_SLICES 32

//inputNHWC_1x480x480x32_outputNHWC_1x480x480x128_kernelHW_1x1_strideHW_1x1_padTopBottomLeftRight_0x0x0x0_dilationHW_1x1
//#define CI 32
//#define IH 480
//#define IW 480
//#define CO 128
//#define OH 480
//#define OW 480
//#define KH 1
//#define KW 1
//#define strideH 1
//#define strideW 1
//#define padTop 0
//#define padLeft 0
//#define CI_SLICES 8
//#define CO_SLICES 32

//inputNHWC_1x60x60x512_outputNHWC_1x60x60x512_kernelHW_1x1_strideHW_1x1_padTopBottomLeftRight_0x0x0x0_dilationHW_1x1
//#define CI 512
//#define IH 60
//#define IW 60
//#define CO 512
//#define OH 60
//#define OW 60
//#define KH 1
//#define KW 1
//#define strideH 1
//#define strideW 1
//#define padTop 0
//#define padLeft 0
//#define CI_SLICES 128
//#define CO_SLICES 128

//#define single_thread_2xFLT4
//#define __global
//#pragma OPENCL EXTENSION cl_arm_printf : enable
//#pragma OPENCL EXTENSION cl_khr_3d_image_writes : enable
#pragma OPENCL EXTENSION cl_khr_fp16 : enable

__constant sampler_t
smp_zero = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;


__kernel void Convolution(__read_only image2d_t input,
                          __write_only image2d_t output,
                          __global FLT4 *weight,
                          __global FLT4 *bias,
                          const int4 input_shape,   // NHWC
                          const int4 output_shape,  // NHWC
                          const int4 kernel_stride, // kernelHW_strideHW
                          const int4 pad)           // top bottom left right)
{
    int oh = get_global_id(1);  // [0, OH)
#ifndef single_thread_2xFLT4
    int ow = get_global_id(0);  // [0, OW)
#else
    int ow = 2 * get_global_id(0);  // [0, OW)
#endif
    int co_slice = get_global_id(2); // [0, UP_DIV(CO, CO_TILE) )

    int CI_SLICES = input_shape.w, IH = input_shape.y, IW = input_shape.z;
    int CO_SLICES = output_shape.w, OH = output_shape.y, OW = output_shape.z;
    int KH = kernel_stride.x, KW = kernel_stride.y;
    int strideH = kernel_stride.z, strideW = kernel_stride.w;
    int padTop = pad.x, padLeft = pad.z;

    if (oh >= OH || ow >= OW || co_slice >= CO_SLICES)
        return;

    int write_second = 1;
    if (ow + 1 >= OW)
        write_second = 0;

    FLT4 out0_c4 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
    __global FLT4 *w0_ic1_oc4 = weight + co_slice * KH * KW * CI_SLICES * CI_TILE;
#ifdef single_thread_2xFLT4
    FLT4 out1_c4 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
#endif

    for (int kh = 0; kh < KH; ++kh)
    {
        int ih = kh + oh * strideH - padTop;
        for (int kw = 0; kw < KW; ++kw)
        {
#ifndef single_thread_2xFLT4
            int iw = kw + ow * strideW - padLeft;
            if (ih >= 0 && ih < IH && iw >= 0 && iw < IW)
            {
                for (int ci_slice = 0; ci_slice < CI_SLICES; ci_slice++)
                {
//                    FLT4 in_c4 = READ_IMAGE(input, smp_zero, (int2)(iw * CI_SLICES + ci_slice, ih)); // NHWC4: H WC
                    FLT4 in_c4 = READ_IMAGE(input, smp_zero, (int2)(iw, ci_slice * IH + ih)); // NC4HW4: CH W
//                    FLT4 in_c4 = READ_IMAGE(input, smp_zero, (int2)(iw, ih * CI_SLICES + ci_slice)); // NHC4W4: HC W
                    out0_c4 += w0_ic1_oc4[0] * in_c4.x;
                    out0_c4 += w0_ic1_oc4[1] * in_c4.y;
                    out0_c4 += w0_ic1_oc4[2] * in_c4.z;
                    out0_c4 += w0_ic1_oc4[3] * in_c4.w;
                    w0_ic1_oc4 += CI_TILE;
                }
            }
            else
            {
                w0_ic1_oc4 += CI_SLICES * CI_TILE;
            }
#else
            // 可以做很大的if else区分write_secondg
            if (ih >= 0 && ih < IH)
            {
                int iw0 = kw + (ow + 0) * strideW - padLeft;
                if (write_second)
                {
                    int iw1 = kw + (ow + 1) * strideW - padLeft;
                    for (int ci_slice = 0; ci_slice < CI_SLICES; ci_slice++)
                    {
                        FLT4 in0_c4 = READ_IMAGE(input, smp_zero, (int2)(iw0, ci_slice * IH + ih)); // NC4HW4: CH W
                        out0_c4 += w0_ic1_oc4[0] * in0_c4.x;
                        out0_c4 += w0_ic1_oc4[1] * in0_c4.y;
                        out0_c4 += w0_ic1_oc4[2] * in0_c4.z;
                        out0_c4 += w0_ic1_oc4[3] * in0_c4.w;
                        FLT4 in1_c4 = READ_IMAGE(input, smp_zero, (int2)(iw1, ci_slice * IH + ih)); // NC4HW4: CH W
                        out1_c4 += w0_ic1_oc4[0] * in1_c4.x;
                        out1_c4 += w0_ic1_oc4[1] * in1_c4.y;
                        out1_c4 += w0_ic1_oc4[2] * in1_c4.z;
                        out1_c4 += w0_ic1_oc4[3] * in1_c4.w;
                        w0_ic1_oc4 += CI_TILE;
                    }
                }
                else
                {
                    for (int ci_slice = 0; ci_slice < CI_SLICES; ci_slice++)
                    {
                        FLT4 in0_c4 = READ_IMAGE(input, smp_zero, (int2)(iw0, ci_slice * IH + ih)); // NC4HW4: CH W
                        out0_c4 += w0_ic1_oc4[0] * in0_c4.x;
                        out0_c4 += w0_ic1_oc4[1] * in0_c4.y;
                        out0_c4 += w0_ic1_oc4[2] * in0_c4.z;
                        out0_c4 += w0_ic1_oc4[3] * in0_c4.w;
                        w0_ic1_oc4 += CI_TILE;
                    }
                }
            }
            else
            {
                w0_ic1_oc4 += CI_SLICES * CI_TILE;
            }
#endif
        }
    }
#ifndef single_thread_2xFLT4
    FLT4 out0_c4_bias = out0_c4 + bias[co_slice];
//    WRITE_IMAGE(output, (int2)(ow * CO_SLICES + co_slice, oh), out0_c4_bias);// NHWC4: H WC
    WRITE_IMAGE(output, (int2)(ow, co_slice * OH + oh), out0_c4_bias);// NC4HW4: CH W
//    WRITE_IMAGE(output, (int2)(ow, oh * CO_SLICES + co_slice), out0_c4_bias);// NHC4W4: HC W
#else
    FLT4 out0_c4_bias = out0_c4 + bias[co_slice];
    WRITE_IMAGE(output, (int2)((ow + 0), co_slice * OH + oh), out0_c4_bias);// NC4HW4: CH W
    if (write_second)
    {
        FLT4 out1_c4_bias = out1_c4 + bias[co_slice];
        WRITE_IMAGE(output, (int2)((ow + 1), co_slice * OH + oh), out1_c4_bias);// NC4HW4: CH W
    }
#endif
}

//__kernel void convolution(__global FLT4 *input,
//                          __global FLT4 *weight,
//                          __global FLT4 *bias,
//                          __global FLT4 *output,
//                          const int4 input_shape,   // NHWC
//                          const int4 output_shape,  // NHWC
//                          const int4 kernel_stride, // kernelHW_strideHW
//                          const int4 pad)           // top bottom left right
//{
//    int oh = get_global_id(0);  // [0, OH)
//    int ow = get_global_id(1);  // [0, OW)
//    int co_slice = get_global_id(2); // [0, UP_DIV(CO, CO_TILE) )
//
//    int CI = input_shape.w, IH = input_shape.y, IW = input_shape.z;
//    int CO = output_shape.w, OH = output_shape.y, OW = output_shape.z;
//    int CI_SLICES = UP_DIV(CI, CI_TILE);
//    int CO_SLICES = UP_DIV(CO, CO_TILE);
//    int KH = kernel_stride.x, KW = kernel_stride.y;
//    int strideH = kernel_stride.z, strideW = kernel_stride.w;
//    int padTop = pad.x, padLeft = pad.z;
//
//    if (oh >= OH || ow >= OW || co_slice >= CO_SLICES)
//        return;
//    FLT4 out0_c4 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
//    __global FLT4 *w0_ic1_oc4 = weight + co_slice * KH * KW * CI_SLICES * CI_TILE;
////    FLT4 out0_c4 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
////    FLT4 out1_c4 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
////    __global FLT4 *w0_ic1_oc4 = weight + (2 * co_slice + 0) * KH * KW * CI_SLICES * CI_TILE;
////    __global FLT4 *w1_ic1_oc4 = weight + (2 * co_slice + 1) * KH * KW * CI_SLICES * CI_TILE;
//    for (int kh = 0; kh < KH; ++kh)
//    {
//        int ih = kh + oh * strideH - padTop;
//        for (int kw = 0; kw < KW; ++kw)
//        {
//            int iw = kw + ow * strideW - padLeft;
//            if (ih >= 0 && ih < IH && iw >= 0 && iw < IW)
//            {
//                for (int ci_slice = 0; ci_slice < CI_SLICES; ci_slice++)
//                {
//                    FLT4 in_c4 = input[ih * IW * CI_SLICES + iw * CI_SLICES + ci_slice];
//                    out0_c4 += w0_ic1_oc4[0] * in_c4.x;
//                    out0_c4 += w0_ic1_oc4[1] * in_c4.y;
//                    out0_c4 += w0_ic1_oc4[2] * in_c4.z;
//                    out0_c4 += w0_ic1_oc4[3] * in_c4.w;
//                    w0_ic1_oc4 += 4;
//
////                    out1_c4 += w1_ic1_oc4[0] * in_c4.x;
////                    out1_c4 += w1_ic1_oc4[1] * in_c4.y;
////                    out1_c4 += w1_ic1_oc4[2] * in_c4.z;
////                    out1_c4 += w1_ic1_oc4[3] * in_c4.w;
////                    w1_ic1_oc4 += 4;
//                }
//            }
//            else
//            {
//                w0_ic1_oc4 += 4 * CI_SLICES;
////                w1_ic1_oc4 += 4 * CI_SLICES;
//            }
//        }
//    }
//    output[oh * OW * CO_SLICES + ow * CO_SLICES + co_slice] = out0_c4 + bias[co_slice];
////    output[oh * OW * CO_SLICES + ow * CO_SLICES + 2 * co_slice + 0] = out0_c4 + bias[2 * co_slice + 0];
////    output[oh * OW * CO_SLICES + ow * CO_SLICES + 2 * co_slice + 1] = out1_c4 + bias[2 * co_slice + 1];
//    return;
//
//    if (oh >= OH || ow >= OW || 2 * co_slice >= CO_SLICES)
//        return;
//    if (2 * co_slice + 1 >= CO_SLICES)
//    {
//        // printf("A %lu %lu %lu\n", get_global_id(0), get_global_id(1), get_global_id(2));
//        FLT4 out0_c4 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
//        __global FLT4 *w0_ic1_oc4 = weight + (2 * co_slice + 0) * KH * KW * CI_SLICES * CI_TILE;
//        for (int kh = 0; kh < KH; ++kh)
//        {
//            int ih = kh + oh * strideH - padTop;
//            for (int kw = 0; kw < KW; ++kw)
//            {
//                int iw = kw + ow * strideW - padLeft;
//                if (ih >= 0 && ih < IH && iw >= 0 && iw < IW)
//                {
//                    for (int ci_slice = 0; ci_slice < CI_SLICES; ci_slice++)
//                    {
//                        FLT4 in_c4 = input[ih * IW * CI_SLICES + iw * CI_SLICES + ci_slice];
//                        out0_c4 += w0_ic1_oc4[0] * in_c4.x;
//                        out0_c4 += w0_ic1_oc4[1] * in_c4.y;
//                        out0_c4 += w0_ic1_oc4[2] * in_c4.z;
//                        out0_c4 += w0_ic1_oc4[3] * in_c4.w;
//                        w0_ic1_oc4 += 4;
//                    }
//                }
//                else
//                {
//                    w0_ic1_oc4 += 4 * CI_SLICES;
//                }
//            }
//        }
//        output[oh * OW * CO_SLICES + ow * CO_SLICES + 2 * co_slice + 0] = out0_c4 + bias[2 * co_slice + 0];
//    }
//    else
//    {
//        // printf("B %lu %lu %lu\n", get_global_id(0), get_global_id(1), get_global_id(2));
//        FLT4 out0_c4 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
//        FLT4 out1_c4 = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
//        __global FLT4 *w0_ic1_oc4 = weight + (2 * co_slice + 0) * KH * KW * CI_SLICES * CI_TILE;
//        __global FLT4 *w1_ic1_oc4 = weight + (2 * co_slice + 1) * KH * KW * CI_SLICES * CI_TILE;
//        for (int kh = 0; kh < KH; ++kh)
//        {
//            int ih = kh + oh * strideH - padTop;
//            for (int kw = 0; kw < KW; ++kw)
//            {
//                int iw = kw + ow * strideW - padLeft;
//                if (ih >= 0 && ih < IH && iw >= 0 && iw < IW)
//                {
//                    int idx = ih * IW * CI_SLICES + iw * CI_SLICES;
//                    for (int ci_slice = 0; ci_slice < CI_SLICES; ci_slice++)
//                    {
//                        FLT4 in_c4 = input[idx + ci_slice];
//
//                        out0_c4 += w0_ic1_oc4[0] * in_c4.x;
//                        out0_c4 += w0_ic1_oc4[1] * in_c4.y;
//                        out0_c4 += w0_ic1_oc4[2] * in_c4.z;
//                        out0_c4 += w0_ic1_oc4[3] * in_c4.w;
//                        w0_ic1_oc4 += 4;
//
//                        out1_c4 += w1_ic1_oc4[0] * in_c4.x;
//                        out1_c4 += w1_ic1_oc4[1] * in_c4.y;
//                        out1_c4 += w1_ic1_oc4[2] * in_c4.z;
//                        out1_c4 += w1_ic1_oc4[3] * in_c4.w;
//                        w1_ic1_oc4 += 4;
//                    }
//                }
//                else
//                {
//                    w0_ic1_oc4 += 4 * CI_SLICES;
//                    w1_ic1_oc4 += 4 * CI_SLICES;
//                }
//            }
//        }
//        output[oh * OW * CO_SLICES + ow * CO_SLICES + 2 * co_slice + 0] = out0_c4 + bias[2 * co_slice + 0];
//        output[oh * OW * CO_SLICES + ow * CO_SLICES + 2 * co_slice + 1] = out1_c4 + bias[2 * co_slice + 1];
//    }
//}